# build reaarrange.py can interactively rearrange the code to make it more readable.

from PIL import Image
import copy
import prior
import random
import time

from ai2thor.controller import Controller
from ai2thor.platform import CloudRendering

def observation(controller): 
    """Return the current (RGB, depth, Optional[instance masks]) frames.

    # Returns
    A tuple containing a
    * RGB frame is of shape (height)x(width)x3 with integer entries in [0:255].
    * depth frame is of shape (height)x(width) with unscaled entries representing the
        meter distance from the agent to the pixel. This will be `None` if the controller_kwargs
        passed to the initializer did not specify that depth images should be returned by AI2-THOR.
    """
    rgb = controller.last_event.frame
    depth = (
        controller.last_event.depth_frame
        if hasattr(controller.last_event, "depth_frame")
        else None
    )
    rgb = Image.fromarray(rgb)
    if depth is not None:
        depth = Image.fromarray((depth / 10 * 255).astype("uint8"), mode="L")
    return rgb, depth

def choose_isPickup(objects):
    """Choose one object from objects that is interactable."""
    interactable_objects = [o for o in objects if o["pickupable"]]
    # print("objects",objects)
    # for interactable_object in interactable_objects:
    #     print("interactable_objects", interactable_object['name'])
    return random.choice(interactable_objects)

def get_top_down_frame(controller):
    # Setup the top-down camera
    event = controller.step(action="GetMapViewCameraProperties", raise_for_failure=True)
    pose = copy.deepcopy(event.metadata["actionReturn"])

    bounds = event.metadata["sceneBounds"]["size"]
    max_bound = max(bounds["x"], bounds["z"])

    pose["fieldOfView"] = 50
    pose["position"]["y"] += 1.1 * max_bound
    pose["orthographic"] = False
    pose["farClippingPlane"] = 50
    del pose["orthographicSize"]

    # add the camera to the scene
    event = controller.step(
        action="AddThirdPartyCamera",
        **pose,
        skyboxColor="white",
        raise_for_failure=True,
    )
    top_down_frame = event.third_party_camera_frames[-1]
    return Image.fromarray(top_down_frame)


def get_pickupable_objects(controller):

    event = controller.last_event
    return event.metadata["arm"]["pickupableObjects"]

dataset = prior.load_dataset("procthor-10k")

controller = Controller(
    agentMode="arm",
    massThreshold=None,
    # scene="FloorPlan203",
    scene=dataset['train'][3],
    visibilityDistance=1,
    gridSize=0.25,
    renderDepthImage=False,
    renderInstanceSegmentation=False,
    width=300,
    height=300,
    fieldOfView=60,
    platform=CloudRendering,
)


# get the object reachable positions

reachable_positions = controller.step(
    action="GetReachablePositions",
).metadata["actionReturn"]

# get object ids

objects = controller.last_event.metadata["objects"]
# print("objects",objects)
target_object = choose_isPickup(objects)      # random choose one object id

print("target_name :",target_object["name"])
print("==="*10)
# get the target object interactable positions  
target_object_interactable_positions =controller.step(
    action="GetInteractablePoses",
    objectId=target_object['objectId'],
    positions=reachable_positions,
).metadata["actionReturn"]


#return one of the interactable positions
target_object_interactable_position = random.choice(target_object_interactable_positions)
print(target_object_interactable_position)

# to Teleport
"""
{'x': 5.5, 'y': 0.9009997844696045, 'z': 2.25, 'rotation': 180.0, 'standing': True, 'horizon': 30.0}

"""
controller.step(
    "Teleport",
    position = dict(x=target_object_interactable_position["x"], y=target_object_interactable_position["y"], z= target_object_interactable_position["z"]), 
    rotation=target_object_interactable_position["rotation"],
    horizon=target_object_interactable_position["horizon"],
    standing=True
)

controller.step(
    action="SetHandSphereRadius",
    radius=0.5,
)

event =controller.step(
    action="PickupObject",
    objectIdCandidates=[target_object["name"]]
)

print(event)

time.sleep(10)
print(event.metadata["lastActionSuccess"])
print("==="*10)

controller.step(action="Pass")  # Ensure the environment is updated before capturing

rgb = controller.last_event.frame
pick_up_frame = Image.fromarray(rgb)
pick_up_frame.save("pick_up_frame.png")


get_top_down_frame(controller).save("top_down_frame.png")




